#include <rclcpp/rclcpp.hpp>
#include <mavros_msgs/msg/state.hpp>
#include <sensor_msgs/msg/nav_sat_fix.hpp>
#include <geometry_msgs/msg/pose_stamped.hpp>
#include <geographic_msgs/msg/geo_pose_stamped.hpp>
#include <mavros_msgs/srv/set_mode.hpp>
#include <mavros_msgs/srv/command_bool.hpp>
#include <GeographicLib/LocalCartesian.hpp>

using namespace std::chrono_literals;

// 坐标转换函数
void wgs84_to_xyz(double lat, double lon, double alt, 
                  double lat_ref, double lon_ref, double alt_ref,
                  double& x, double& y, double& z) {
    GeographicLib::LocalCartesian local(lat_ref, lon_ref, alt_ref);
    local.Forward(lat, lon, alt, x, y, z);
}

void xyz_to_wgs84(double x, double y, double z, 
                  double lat_ref, double lon_ref, double alt_ref,
                  double& lat, double& lon, double& alt) {
    GeographicLib::LocalCartesian local(lat_ref, lon_ref, alt_ref);
    local.Reverse(x, y, z, lat, lon, alt);
}

class OffboardControl : public rclcpp::Node {
public:
    OffboardControl() : Node("offboard_control_node") {
        // 初始化订阅者、发布者和服务客户端
        state_sub_ = this->create_subscription<mavros_msgs::msg::State>(
            "/mavros/state", rclcpp::QoS(20).best_effort(), std::bind(&OffboardControl::state_callback, this, std::placeholders::_1));
        global_pos_sub_ = this->create_subscription<sensor_msgs::msg::NavSatFix>(
            "/mavros/global_position/global", rclcpp::QoS(20).best_effort(), std::bind(&OffboardControl::global_pos_callback, this, std::placeholders::_1));
        global_pos_pub_ = this->create_publisher<geographic_msgs::msg::GeoPoseStamped>(
            "/mavros/setpoint_position/global", 10);
        local_pos_pub_ = this->create_publisher<geometry_msgs::msg::PoseStamped>(
            "/mavros/setpoint_position/local", 10);

        set_mode_client_ = this->create_client<mavros_msgs::srv::SetMode>("/mavros/set_mode");
        arming_client_ = this->create_client<mavros_msgs::srv::CommandBool>("/mavros/cmd/arming");

        // 初始化定时器
        timer_ = this->create_wall_timer(
            50ms, std::bind(&OffboardControl::timer_callback, this));

        RCLCPP_INFO(this->get_logger(), "OffboardControl Node initialized.");
    }

private:
    float start = 0;
    int Cnt = 0;
    void state_callback(const mavros_msgs::msg::State::SharedPtr msg) {
        current_state_ = *msg;
        if(current_state_.armed)
        {
            start = 1;
        }
        RCLCPP_INFO(this->get_logger(), "Armed State is: %f",start);
    }

    void global_pos_callback(const sensor_msgs::msg::NavSatFix::SharedPtr msg) {
        current_global_pose_ = *msg;

        if (!start) {
            // 设置 Home 点为参考点
            home_lat_ = msg->latitude;
            home_lon_ = msg->longitude;
            home_alt_ = msg->altitude;
            is_home_set_ = true;

            RCLCPP_INFO(this->get_logger(), "Home position set: lat=%.6f, lon=%.6f, alt=%.2f",
                        home_lat_, home_lon_, home_alt_);


            // double home_x_, home_y_, home_z_;
            // wgs84_to_xyz(double lat, double lon, double alt, 
            //       home_lat_, home_lon_, home_alt_,
            //       home_x_, home_y_, home_z_) 

            // 设置起飞目标
            // target_x_ = home_x_ + 0.0; // 局部坐标原点
            // target_y_ = home_y_ + 0.0;
            // target_z_ = home_z_ + 2.0; // 起飞目标高度（相对 Home 点 +2 米）
        }
    }

    void timer_callback() {
        if (!current_state_.connected) {
            RCLCPP_WARN(this->get_logger(), "Waiting for FCU connection...");
            return;
        }

        if (initial_setup_) {
            perform_initial_setup();
            return;
        }

        handle_offboard_mode();
        handle_arming();

        if (is_home_set_) {
            if(Cnt<100)
            {
                if(start)
                {
                    Cnt++;
                    geographic_msgs::msg::GeoPoseStamped target_pose;
                    target_pose.header.stamp = this->now();

                    // 使用全球坐标设置目标位置
                    target_pose.pose.position.latitude = home_lat_;
                    target_pose.pose.position.longitude = home_lon_;
                    target_pose.pose.position.altitude = home_alt_;
                    global_pos_pub_->publish(target_pose);
                    RCLCPP_INFO(this->get_logger(), "Cnt: %d", Cnt); 
                }
            }
            else
            {
                publish_local_target_pose();
            }
            
        } else {
            RCLCPP_INFO(this->get_logger(), "Waiting for home position to be set...");
        }
    }

    void perform_initial_setup() {
        RCLCPP_INFO(this->get_logger(), "Performing initial setup...");
        for (int i = 0; i < 100 && rclcpp::ok(); ++i) {
            geographic_msgs::msg::GeoPoseStamped target_pose;
            target_pose.header.stamp = this->now();

            // 使用全球坐标设置目标位置
            target_pose.pose.position.latitude = home_lat_;
            target_pose.pose.position.longitude = home_lon_;
            target_pose.pose.position.altitude = home_alt_;
            global_pos_pub_->publish(target_pose);
            RCLCPP_INFO(this->get_logger(), "Cnt:%d", Cnt);
            rclcpp::sleep_for(50ms);
        }
        initial_setup_ = false;
        last_request_ = this->get_clock()->now();
        RCLCPP_INFO(this->get_logger(), "Initial setup complete.");
    }

    void handle_offboard_mode() {
        if (current_state_.mode != "OFFBOARD" &&
            (this->get_clock()->now() - last_request_ > rclcpp::Duration(0.2))) {
            auto request = std::make_shared<mavros_msgs::srv::SetMode::Request>();
            request->custom_mode = "OFFBOARD";

            if (set_mode_client_->wait_for_service(100ms)) {
                auto result = set_mode_client_->async_send_request(request);
                if (result.wait_for(10ms) == std::future_status::ready && result.get()->mode_sent) {
                    RCLCPP_INFO(this->get_logger(), "Offboard mode enabled.");
                } else {
                    RCLCPP_WARN(this->get_logger(), "Failed to set Offboard mode.");
                }
            }
            last_request_ = this->get_clock()->now();
        }
    }

    void handle_arming() {
        if (!current_state_.armed &&
            (this->get_clock()->now() - last_request_ > rclcpp::Duration(0.0))) {
            auto request = std::make_shared<mavros_msgs::srv::CommandBool::Request>();
            request->value = true;

            if (arming_client_->wait_for_service(10ms)) {
                auto result = arming_client_->async_send_request(request);
                if (result.wait_for(10ms) == std::future_status::ready && result.get()->success) {
                    RCLCPP_INFO(this->get_logger(), "Vehicle armed.");
                    start = 1;
                } else {
                    RCLCPP_WARN(this->get_logger(), "Failed to arm the vehicle.");
                }
            }
            last_request_ = this->get_clock()->now();
        }
    }

    void publish_local_target_pose() {
        geometry_msgs::msg::PoseStamped target_pose;
        target_pose.header.stamp = this->now();
        target_x_ = 0;
        target_y_ = 0;
        target_z_ = 2;
        // NED
        // 使用局部坐标设置目标位置
        target_pose.pose.position.x = target_x_;
        target_pose.pose.position.y = target_y_;
        target_pose.pose.position.z = target_z_;
        // target_pose.pose.orientation.x = NAN;
        // target_pose.pose.orientation.y = NAN;
        // target_pose.pose.orientation.z = NAN;
        // target_pose.pose.orientation.w = NAN;

        local_pos_pub_->publish(target_pose);
        RCLCPP_INFO(this->get_logger(), "Publishing target pose (local): x=%.2f, y=%.2f, z=%.2f",
                    target_x_, target_y_, target_z_);
    }

    // ROS组件
    rclcpp::Subscription<mavros_msgs::msg::State>::SharedPtr state_sub_;
    rclcpp::Subscription<sensor_msgs::msg::NavSatFix>::SharedPtr global_pos_sub_;
    rclcpp::Publisher<geographic_msgs::msg::GeoPoseStamped>::SharedPtr global_pos_pub_;
    rclcpp::Publisher<geometry_msgs::msg::PoseStamped>::SharedPtr local_pos_pub_;
    rclcpp::Client<mavros_msgs::srv::SetMode>::SharedPtr set_mode_client_;
    rclcpp::Client<mavros_msgs::srv::CommandBool>::SharedPtr arming_client_;
    rclcpp::TimerBase::SharedPtr timer_;

    // 无人机状态和控制变量
    mavros_msgs::msg::State current_state_;
    sensor_msgs::msg::NavSatFix current_global_pose_;
    rclcpp::Time last_request_;
    bool initial_setup_ = true;

    // Home点和目标点
    double home_lat_, home_lon_, home_alt_;
    double target_x_, target_y_, target_z_;
    bool is_home_set_ = false;
};

int main(int argc, char **argv) {
    rclcpp::init(argc, argv);
    auto node = std::make_shared<OffboardControl>();
    rclcpp::spin(node);
    rclcpp::shutdown();
    return 0;
}
